#include "parser.h"
#include <stdexcept>

Parser::Parser(const std::vector<Token>& tokens) 
    : tokens(tokens), current(0) {}

Token Parser::currentToken() {
    if (current >= tokens.size()) {
        return Token(TokenType::END_OF_FILE, "", 0, 0);
    }
    return tokens[current];
}

Token Parser::peekToken() {
    if (current + 1 >= tokens.size()) {
        return Token(TokenType::END_OF_FILE, "", 0, 0);
    }
    return tokens[current + 1];
}

void Parser::advance() {
    if (current < tokens.size()) {
        current++;
    }
}

bool Parser::match(TokenType type) {
    if (isAtEnd()) return false;
    if (currentToken().type == type) {
        advance();
        return true;
    }
    return false;
}

bool Parser::isAtEnd() {
    return currentToken().type == TokenType::END_OF_FILE;
}

std::unique_ptr<Program> Parser::parse() {
    std::vector<std::unique_ptr<Declaration>> declarations;
    
    while (!isAtEnd()) {
        try {
            declarations.push_back(parseDeclaration());
        } catch (const std::runtime_error& e) {
            synchronize();
        }
    }
    
    return std::make_unique<Program>(std::move(declarations));
}

std::unique_ptr<Declaration> Parser::parseDeclaration() {
    if (currentToken().type == TokenType::INT || 
        currentToken().type == TokenType::CHAR ||
        currentToken().type == TokenType::FLOAT ||
        currentToken().type == TokenType::DOUBLE ||
        currentToken().type == TokenType::VOID) {
        
        TokenType type = currentToken().type;
        advance();
        
        if (peekToken().type == TokenType::LEFT_PAREN) {
            return parseFunctionDeclaration();
        } else {
            return parseVariableDeclaration();
        }
    }
    
    throw std::runtime_error("Expected declaration");
}

std::unique_ptr<Declaration> Parser::parseVariableDeclaration() {
    TokenType type = currentToken().type;
    advance();
    
    if (currentToken().type != TokenType::IDENTIFIER) {
        throw std::runtime_error("Expected identifier");
    }
    
    std::string name = currentToken().value;
    advance();
    
    std::unique_ptr<Expression> initializer = nullptr;
    if (match(TokenType::ASSIGN)) {
        initializer = parseExpression();
    }
    
    if (!match(TokenType::SEMICOLON)) {
        throw std::runtime_error("Expected ';' after variable declaration");
    }
    
    return std::make_unique<VariableDeclaration>(type, name, std::move(initializer));
}

std::unique_ptr<Declaration> Parser::parseFunctionDeclaration() {
    TokenType returnType = currentToken().type;
    advance();
    
    if (currentToken().type != TokenType::IDENTIFIER) {
        throw std::runtime_error("Expected function name");
    }
    
    std::string name = currentToken().value;
    advance();
    
    if (!match(TokenType::LEFT_PAREN)) {
        throw std::runtime_error("Expected '(' after function name");
    }
    
    std::vector<std::pair<TokenType, std::string>> parameters;
    if (!match(TokenType::RIGHT_PAREN)) {
        do {
            if (currentToken().type == TokenType::INT ||
                currentToken().type == TokenType::CHAR ||
                currentToken().type == TokenType::FLOAT ||
                currentToken().type == TokenType::DOUBLE) {
                
                TokenType paramType = currentToken().type;
                advance();
                
                if (currentToken().type != TokenType::IDENTIFIER) {
                    throw std::runtime_error("Expected parameter name");
                }
                
                parameters.emplace_back(paramType, currentToken().value);
                advance();
            } else {
                throw std::runtime_error("Expected parameter type");
            }
        } while (match(TokenType::COMMA));
        
        if (!match(TokenType::RIGHT_PAREN)) {
            throw std::runtime_error("Expected ')' after parameters");
        }
    }
    
    auto body = parseBlockStatement();
    return std::make_unique<FunctionDeclaration>(returnType, name, std::move(parameters), std::move(body));
}

std::unique_ptr<Statement> Parser::parseStatement() {
    if (currentToken().type == TokenType::RETURN) {
        return parseReturnStatement();
    } else if (currentToken().type == TokenType::IF) {
        return parseIfStatement();
    } else if (currentToken().type == TokenType::WHILE) {
        return parseWhileStatement();
    } else if (currentToken().type == TokenType::LEFT_BRACE) {
        return parseBlockStatement();
    } else {
        return parseExpressionStatement();
    }
}

std::unique_ptr<Statement> Parser::parseExpressionStatement() {
    auto expr = parseExpression();
    if (!match(TokenType::SEMICOLON)) {
        throw std::runtime_error("Expected ';' after expression");
    }
    return std::make_unique<ExpressionStatement>(std::move(expr));
}

std::unique_ptr<Statement> Parser::parseReturnStatement() {
    advance(); // consume 'return'
    
    std::unique_ptr<Expression> expr = nullptr;
    if (currentToken().type != TokenType::SEMICOLON) {
        expr = parseExpression();
    }
    
    if (!match(TokenType::SEMICOLON)) {
        throw std::runtime_error("Expected ';' after return statement");
    }
    
    return std::make_unique<ReturnStatement>(std::move(expr));
}

std::unique_ptr<Statement> Parser::parseIfStatement() {
    advance(); // consume 'if'
    
    if (!match(TokenType::LEFT_PAREN)) {
        throw std::runtime_error("Expected '(' after if");
    }
    
    auto condition = parseExpression();
    
    if (!match(TokenType::RIGHT_PAREN)) {
        throw std::runtime_error("Expected ')' after if condition");
    }
    
    auto thenBranch = parseStatement();
    std::unique_ptr<Statement> elseBranch = nullptr;
    
    if (match(TokenType::ELSE)) {
        elseBranch = parseStatement();
    }
    
    return std::make_unique<IfStatement>(std::move(condition), std::move(thenBranch), std::move(elseBranch));
}

std::unique_ptr<Statement> Parser::parseWhileStatement() {
    advance(); // consume 'while'
    
    if (!match(TokenType::LEFT_PAREN)) {
        throw std::runtime_error("Expected '(' after while");
    }
    
    auto condition = parseExpression();
    
    if (!match(TokenType::RIGHT_PAREN)) {
        throw std::runtime_error("Expected ')' after while condition");
    }
    
    auto body = parseStatement();
    
    return std::make_unique<WhileStatement>(std::move(condition), std::move(body));
}

std::unique_ptr<Statement> Parser::parseBlockStatement() {
    if (!match(TokenType::LEFT_BRACE)) {
        throw std::runtime_error("Expected '{'");
    }
    
    std::vector<std::unique_ptr<Statement>> statements;
    
    while (currentToken().type != TokenType::RIGHT_BRACE && !isAtEnd()) {
        statements.push_back(parseStatement());
    }
    
    if (!match(TokenType::RIGHT_BRACE)) {
        throw std::runtime_error("Expected '}' after block");
    }
    
    return std::make_unique<BlockStatement>(std::move(statements));
}

std::unique_ptr<Expression> Parser::parseExpression() {
    return parseAssignment();
}

std::unique_ptr<Expression> Parser::parseAssignment() {
    auto expr = parseEquality();
    
    if (match(TokenType::ASSIGN)) {
        auto value = parseAssignment();
        // For now, we'll treat assignment as a binary expression
        return std::make_unique<BinaryExpression>(std::move(expr), TokenType::ASSIGN, std::move(value));
    }
    
    return expr;
}

std::unique_ptr<Expression> Parser::parseEquality() {
    auto expr = parseComparison();
    
    while (match(TokenType::EQUAL) || match(TokenType::NOT_EQUAL)) {
        TokenType operatorType = tokens[current - 1].type;
        auto right = parseComparison();
        expr = std::make_unique<BinaryExpression>(std::move(expr), operatorType, std::move(right));
    }
    
    return expr;
}

std::unique_ptr<Expression> Parser::parseComparison() {
    auto expr = parseTerm();
    
    while (match(TokenType::GREATER) || match(TokenType::GREATER_EQUAL) ||
           match(TokenType::LESS) || match(TokenType::LESS_EQUAL)) {
        TokenType operatorType = tokens[current - 1].type;
        auto right = parseTerm();
        expr = std::make_unique<BinaryExpression>(std::move(expr), operatorType, std::move(right));
    }
    
    return expr;
}

std::unique_ptr<Expression> Parser::parseTerm() {
    auto expr = parseFactor();
    
    while (match(TokenType::PLUS) || match(TokenType::MINUS)) {
        TokenType operatorType = tokens[current - 1].type;
        auto right = parseFactor();
        expr = std::make_unique<BinaryExpression>(std::move(expr), operatorType, std::move(right));
    }
    
    return expr;
}

std::unique_ptr<Expression> Parser::parseFactor() {
    auto expr = parseUnary();
    
    while (match(TokenType::MULTIPLY) || match(TokenType::DIVIDE) || match(TokenType::MODULO)) {
        TokenType operatorType = tokens[current - 1].type;
        auto right = parseUnary();
        expr = std::make_unique<BinaryExpression>(std::move(expr), operatorType, std::move(right));
    }
    
    return expr;
}

std::unique_ptr<Expression> Parser::parseUnary() {
    if (match(TokenType::NOT) || match(TokenType::MINUS)) {
        TokenType operatorType = tokens[current - 1].type;
        auto right = parseUnary();
        return std::make_unique<UnaryExpression>(operatorType, std::move(right));
    }
    
    return parsePrimary();
}

std::unique_ptr<Expression> Parser::parsePrimary() {
    if (match(TokenType::FALSE)) {
        return std::make_unique<Expression>(TokenType::UNKNOWN, "false");
    }
    
    if (match(TokenType::TRUE)) {
        return std::make_unique<Expression>(TokenType::UNKNOWN, "true");
    }
    
    if (match(TokenType::INTEGER_LITERAL) || match(TokenType::FLOAT_LITERAL) || 
        match(TokenType::CHAR_LITERAL) || match(TokenType::STRING_LITERAL)) {
        TokenType type = tokens[current - 1].type;
        std::string value = tokens[current - 1].value;
        return std::make_unique<Expression>(type, value);
    }
    
    if (match(TokenType::IDENTIFIER)) {
        std::string name = tokens[current - 1].value;
        
        if (match(TokenType::LEFT_PAREN)) {
            std::vector<std::unique_ptr<Expression>> arguments;
            
            if (!match(TokenType::RIGHT_PAREN)) {
                do {
                    arguments.push_back(parseExpression());
                } while (match(TokenType::COMMA));
                
                if (!match(TokenType::RIGHT_PAREN)) {
                    throw std::runtime_error("Expected ')' after arguments");
                }
            }
            
            return std::make_unique<FunctionCall>(name, std::move(arguments));
        }
        
        return std::make_unique<Expression>(TokenType::IDENTIFIER, name);
    }
    
    if (match(TokenType::LEFT_PAREN)) {
        auto expr = parseExpression();
        if (!match(TokenType::RIGHT_PAREN)) {
            throw std::runtime_error("Expected ')' after expression");
        }
        return expr;
    }
    
    throw std::runtime_error("Expected expression");
}

void Parser::synchronize() {
    advance();
    
    while (!isAtEnd()) {
        if (tokens[current - 1].type == TokenType::SEMICOLON) {
            return;
        }
        
        switch (currentToken().type) {
            case TokenType::INT:
            case TokenType::CHAR:
            case TokenType::FLOAT:
            case TokenType::DOUBLE:
            case TokenType::VOID:
            case TokenType::IF:
            case TokenType::WHILE:
            case TokenType::FOR:
            case TokenType::RETURN:
                return;
            default:
                break;
        }
        
        advance();
    }
}
